Skip to content

refactor(triton): simplify JIT backend architecture - #925

Draft
voltjia wants to merge 6 commits into
stack/pr-800-triton-backend-jitfrom
refactor/triton-jit-architecture
Draft

refactor(triton): simplify JIT backend architecture#925
voltjia wants to merge 6 commits into
stack/pr-800-triton-backend-jitfrom
refactor/triton-jit-architecture

Conversation

@voltjia

@voltjia voltjia commented Aug 12, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Move the Triton JIT implementation into infini::ops::triton::jit and use concise names such as Target, DataType, Backend, ScopedDevice, and Kernel.
  • Keep the reusable InfiniRT Runtime/Driver implementation in backend.h, with an inline NVIDIA Backend specialization; remove the separate NVIDIA declaration/instantiation files.
  • Align Target with Triton's GPUTarget semantics (backend, architecture, and warp_size) and pass the runtime device ID separately.
  • Treat KernelMetadata as the local projection needed to load and launch a compiled Triton kernel: name, actual num_warps, shared memory, and unsupported scratch sizes. Keep the binary extension in the artifact envelope rather than kernel metadata.
  • Replace hand-written Python C API object management with pybind11. The only direct C API call left is Py_IsInitialized(), which safely guards GIL acquisition when no interpreter exists.
  • Move cache filesystem operations to Python, replace ad hoc identity helpers with IdentityBuilder, and use value/optional returns instead of bool-plus-output-pointer APIs.
  • Rename the Add Triton source to jit.py, keep pyproject.toml unchanged, and retain config_.cc as the key-function/vtable anchor for polymorphic configs across the core and pybind DSOs.
  • Preserve explicit release-build short circuits for invalid devices, targets, arguments, grids, cache paths, metadata uploads, kernel loads, and auto-tuning failures.

Motivation

This is a stacked cleanup of #800 focused on architecture, naming, modularity,
readability, and maintainability while preserving the template boundary needed
for future Triton backends.

Both stacked branches are rebased onto master at
7f4201e19eddf8c5eaf5cdbdde6ba8ae8329040a. This PR contains one architecture
commit on top of stack/pr-800-triton-backend-jit; the already merged #919 and
#924 changes are inherited from master.

Related: #800, #919

Type of Change

  • refactor - code restructuring without an intended public behavior change
  • Breaking change

Platforms Affected

  • CPU (WITH_CPU)
  • NVIDIA (WITH_NVIDIA)
  • Iluvatar (WITH_ILUVATAR)
  • MetaX (WITH_METAX)
  • Cambricon (WITH_CAMBRICON)
  • Moore (WITH_MOORE)
  • Ascend (WITH_ASCEND)
  • PyTorch C++ bindings (WITH_TORCH)
  • Build system / CMake / CI
  • Python bindings / user-facing API

Smoke Test Result

Focused validation for current commit
1a0197a5920351ca7ae6b24a789f983ff438f18a after the final rebase:

Ruff check and format check on changed Python files: PASS
clang-format 21 dry run on changed C++ files: PASS
focused generator/compiler tests: 42 passed
Python bytecode compilation: PASS
git diff --check and stale-reference audit: PASS

The predecessor architecture was built and exercised on NVIDIA, including
configured Add, auto-tuning, and multi-device smoke. The backend and bridge have
changed since that run, so a full build and GPU smoke will be rerun on the final
SHA before merge rather than attributed to this commit.

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA Yes Focused checks pass Current-SHA build and GPU smoke pending the pre-merge validation run
Iluvatar No N/A
MetaX No N/A
Cambricon No N/A
Moore No N/A
Ascend No N/A

Benchmark / Performance Impact

N/A. This PR does not make a performance claim.

Notes for Reviewers

  • feat(triton): add JIT backend with add operator  #800 comes from a fork, so GitHub cannot use its head directly as an upstream PR base. stack/pr-800-triton-backend-jit mirrors its three commits and is now rebased onto current master; the source PR feat(triton): add JIT backend with add operator  #800 still points to its fork head.
  • Triton's GPUTarget contains backend, architecture, and warp size. Runtime device selection therefore stays outside Target.
  • Triton builds metadata dynamically and loads it as a named tuple. This PR deliberately keeps only the fields consumed by InfiniOps; binary_ext remains an artifact-format property rather than a KernelMetadata field.
  • Python calls with an explicit Triton config intentionally use uncached Make so derived config identity is not erased by the generic operator cache. Compiled kernels and auto-tuning results remain cached by their complete identities. Configless calls continue through Call.
  • pybind11 internally reports Python failures with C++ exceptions. The bridge contains those exceptions at the Python boundary and converts them to empty status values; they do not escape into the runtime/operator layers.
  • A future Triton backend adds a Backend<kDev> specialization backed by the corresponding InfiniRT Runtime/Driver capabilities. Unsupported devices remain invisible to operator implementation discovery through the incomplete primary template.

@voltjia
voltjia force-pushed the refactor/triton-jit-architecture branch from de9694a to bd9f92d Compare August 12, 2026 11:14
@voltjia
voltjia force-pushed the stack/pr-800-triton-backend-jit branch from 0c73a9e to 1c9ec2c Compare August 12, 2026 22:22
@voltjia
voltjia force-pushed the refactor/triton-jit-architecture branch from bd9f92d to 439472d Compare August 12, 2026 22:23
@voltjia
voltjia force-pushed the stack/pr-800-triton-backend-jit branch from 1c9ec2c to e9d98ec Compare August 13, 2026 05:54
@voltjia
voltjia force-pushed the refactor/triton-jit-architecture branch 2 times, most recently from 1a0197a to 3f8f072 Compare August 13, 2026 06:14
@voltjia
voltjia force-pushed the refactor/triton-jit-architecture branch from 3f8f072 to 1a6b83c Compare August 13, 2026 06:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant